#!/usr/bin/env python
#-*- coding:utf-8 -*-

# --------------------------------------------------------
# Faster R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""
调用须知：

1. 该源码:
   用来生成图片的xml

2. 调用前的准备工作

(1). 大前提: 确保CLASSES是你要检测的类别

"""

import matplotlib
matplotlib.use("Agg")
import _init_paths
from fast_rcnn.config import cfg, cfg_from_file, cfg_from_list
from fast_rcnn.test import im_detect
from fast_rcnn.nms_wrapper import nms
from utils.timer import Timer
import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio
import caffe, os, sys, cv2
import argparse
import shutil
from collections import OrderedDict
from pandas import DataFrame
from pascal_voc_io import PascalVocWriter
import time

CLASSES = ('__background__', 'hat', 'man')
dict_list = []
current_time = time.strftime('%Y-%m-%d_%H-%M-%S', time.localtime(time.time()))

def add_boxes(xml, image_name, class_name, dets, inds):
    """
    Draw detected bounding boxes.
    Add the varible 'ax', 'inds'.
    Delete 'im' .
    """
    for i in inds:
        # row = OrderedDict([('ImgName', os.path.splitext(image_name)[0]), ('ClsName', class_name), ('score', dets[i, -1]),
        #                     ('xmin', dets[i, 0]), ('ymin', dets[i, 1]), ('xmax', dets[i, 2]), ('ymax', dets[i, 3])])
        # dict_list.append(row)
        xml.addBndBox(int(dets[i, 0]) + 1, int(dets[i, 1]) + 1, int(dets[i, 2]) + 1, int(dets[i, 3]) + 1, class_name)

def auto_gen_xml(net, image_name):
    """Detect object classes in an image using pre-computed object proposals."""

    ####### ----Newly added----Set the Savepath------ ############
    xml_savepath = os.path.join(cfg.DATA_DIR, "auto_gen_xml", current_time, 'xml')
    img_savepath = os.path.join(cfg.DATA_DIR, "auto_gen_xml", current_time, 'img')
    not_detect_savepath = os.path.join(cfg.DATA_DIR, "auto_gen_xml", current_time, "not_detect")
    if not os.path.isdir(xml_savepath):
        os.makedirs(xml_savepath)
    if not os.path.isdir(img_savepath):
        os.makedirs(img_savepath)
    if not os.path.isdir(not_detect_savepath):
        os.makedirs(not_detect_savepath)

    # Load the image
    im_file = os.path.join(args.imgs_absdir, image_name)
    im = cv2.imread(im_file)

    # Detect all object classes and regress object bounds
    timer = Timer()
    timer.tic()
    scores, boxes = im_detect(net, im)
    timer.toc()
    print(('Detection took {:.3f}s for '
           '{:d} object proposals').format(timer.total_time, boxes.shape[0]))

    # detections for each class
    CONF_THRESH = 0.9
    NMS_THRESH = 0.3

    hrs_Num = 0

    xml = PascalVocWriter('VOC2007', image_name, im.shape)

    for cls_ind, cls in enumerate(CLASSES[1:]):
        cls_ind += 1 # because we skipped background
        cls_boxes = boxes[:, 4*cls_ind:4*(cls_ind + 1)]
        cls_scores = scores[:, cls_ind]
        dets = np.hstack((cls_boxes,
                          cls_scores[:, np.newaxis])).astype(np.float32)
        keep = nms(dets, NMS_THRESH)  #非极大值抑制
        dets = dets[keep, :]
        # ----Newly added----获取概率大于CONF_THRESH的dets第一维度的索引------ #
        inds = np.where(dets[:, -1] >= CONF_THRESH)[0]

        # ----Newly added---对于某类的相似度低于CONF_THRESH的图片, 计数加1----- #
        if len(inds) == 0:
            hrs_Num += 1
            # pics_filename = os.path.join(not_detect_savepath, os.path.splitext(image_name)[0] + '_Not_' + cls + '.jpg')
            # shutil.copy(im_file, pics_filename)
            # 若需要保存既含hat又含man的xml, 当inds为空时, 则返回空值, 结束整个demo方法, 不保存xml
            # 不需要保存既含hat又含man的xml, 当inds为空时, 则不执行本次循环剩下的语句, 直接进行本层循环的下一次循环, 而不会结束循环
            if args.check_hatman:
                return
            else:
                continue

        add_boxes(xml, image_name, cls, dets, inds)

    if hrs_Num == len(CLASSES) - 1:
        # 若CLASSES类别中的所有类的相似度都低于CONF_THRESH, 则将图片COPY到not_detect_savepath
        shutil.copy(im_file, not_detect_savepath)
        return

    xml_filename = os.path.join(xml_savepath, image_name[:-4] + '.xml')
    xml.save(xml_filename)
    shutil.copy(im_file, img_savepath)

def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser(description='auto detection and generating xml')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]',
                        default=0, type=int)
    parser.add_argument('--cpu', dest='cpu_mode',
                        help='Use CPU mode (overrides --gpu)',
                        action='store_true')
    parser.add_argument('--net', dest='caffemodel', help='model to test',
                        default=None, type=str)
    parser.add_argument('--def', dest='prototxt',
                        help='prototxt file defining the network',
                        default=None, type=str)
    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file', default=None, type=str)
    parser.add_argument('--set', dest='set_cfgs',
                        help='set config keys', default=None, nargs=argparse.REMAINDER)
    parser.add_argument('--imgdir', dest='imgs_absdir',
                        help='set images path', default=None, type=str)
    parser.add_argument('--check_hat_man', dest='check_hatman',
                        help='set the mode when needing check only have one hat and one man',
                        action='store_true')

    args = parser.parse_args()

    return args


if __name__ == '__main__':

    cfg.TEST.HAS_RPN = True  # Use RPN for proposals

    args = parse_args()

    if args.cfg_file is not None:
        cfg_from_file(args.cfg_file)
    if args.set_cfgs is not None:
        cfg_from_list(args.set_cfgs)

    caffemodel = args.caffemodel
    prototxt = args.prototxt

    if not os.path.isfile(caffemodel):
        raise IOError(('{:s} not found.\nDid you run ./data/script/'
                       'fetch_faster_rcnn_models.sh?').format(caffemodel))

    if args.cpu_mode:
        caffe.set_mode_cpu()
    else:
        caffe.set_mode_gpu()
        caffe.set_device(args.gpu_id)
        cfg.GPU_ID = args.gpu_id

    net = caffe.Net(prototxt, caffemodel, caffe.TEST)

    print('\n\nLoaded network {:s}'.format(caffemodel))

    # imgpath = os.path.join(args.imgs_path)
    assert args.imgs_absdir, "images absdir doesn't exist, please check it. "
    im_names = os.listdir(args.imgs_absdir)
    im_names.sort()

    for im_name in im_names:
        print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
        print('data--{}'.format(im_name))
        auto_gen_xml(net, im_name)

    # d1 = DataFrame(dict_list)
    # d1 = d1.sort_values(by=['ImgName', 'ClsName', 'score', 'xmin', 'ymin'])
    # csvsavepath1 = os.path.join(cfg.DATA_DIR,'auto_gen_xml', current_time + '.csv')
    # if not os.path.exists(csvsavepath1):
    #     d1.to_csv(csvsavepath1, index=False, header=True)
